#!/usr/bin/env python
"""
Constructs the Hamiltonian H for a transverse field Ising model and
explicitly diagonalizes it.

N - total number of spins in the lattice.


"""


from __future__ import division
import numpy as np
import scipy
import pickle

def get_Sz(i,N):
    Sz = np.zeros(N)
    b_str = np.binary_repr(i, N)
    for k in range(N):
        Sz[N - 1 - k] = np.int(b_str[k])
    Sz = 2 * Sz - 1
    return Sz


def get_state(Sz):
    N = np.size(Sz)
    Sz = (Sz + 1)/2
    i = np.int( np.sum(Sz * (2**np.arange(N))) )
    return i

def get_xyfromi(i,N):
    x = i%N
    y = int(i/N)
    return x,y

def get_ifromxy(x, y, N):
    i = y*N + x
    return i

def get_neigh(i,N):
    """
    The function in which information about the lattice is present
    """
    L = int(np.sqrt(N))
    n = [L-1] + range(L-1)
    w = n
    s = range(1, L) + [0]
    e = s
    neigh = []
    x, y = get_xyfromi(i, L)
    neigh.append(get_ifromxy(e[x], y, L))
    neigh.append(get_ifromxy(w[x], y, L))
    neigh.append(get_ifromxy(x, n[y], L))
    neigh.append(get_ifromxy(x, s[y], L))    
    neigh.append(get_ifromxy(n[x], e[y], L))    
    neigh.append(get_ifromxy(s[x], w[y], L))
    return neigh

def sum_neigh(Sz, i):
    N = np.size(Sz)
    t_sum = np.sum(Sz[get_neigh(i, N)])
    return t_sum
    

def get_diag(i, N):
    """
    This is the term that gets you the ith diagonal term of the
    Hamiltonian.
    """
    Sz = get_Sz(i, N)
    t = 0
    for j in range(N):
        t -= Sz[j] * sum_neigh(Sz, j)
    t *= 0.5
    return t

    

def construct_hamiltonian(N, J, B):
    dim = 2**N
    H = np.zeros([dim,dim])
    for i in range(dim):
        H[i][i] = J * get_diag(i, N)
        Sz = get_Sz(i, N)
        for j in range(N):
            Sz[j] *= -1
            k = get_state(Sz)
            Sz[j] *= -1
            H[i][k] = B
    return H
        

def test_symmetric(M):
    M = np.matrix(M)
    return np.all(M == M.T)



N = 9
B = 0.0
H = construct_hamiltonian(N, 0.1, B)
print test_symmetric(H)
# print H
# print np.sum(H,1)

eig, eigvec = np.linalg.eig(H)
prefix = 'exact_diag_N_%d_B_%0.2f_eig'%(N,B)
fname = prefix + '_values.pick'
pickle.dump(eig,open(fname,"wb"))
fname = prefix + '_vectors.pick'
pickle.dump(eigvec,open(fname,"wb"))

